import os
import django
import sys
from openpyxl import Workbook
from openpyxl.styles import Font, Alignment


# pip install openpyxl

def export_model_to_excel(model, file_path, exclude_fields=None):
    """
    将Django模型数据导出到Excel

    参数:
    - model: 要导出的Django模型类
    - file_path: 导出的Excel文件路径
    - exclude_fields: 要排除的字段列表，默认为None
    """
    # 初始化Workbook
    wb = Workbook()
    ws = wb.active
    ws.title = model.__name__
    
    # 获取模型字段
    exclude_fields = exclude_fields or []
    fields = [field for field in model._meta.get_fields()
              if field.name not in exclude_fields
              and hasattr(field, 'verbose_name')
              and not field.auto_created]  # 排除自动创建的字段（如反向关联）
    
    # 设置表头
    for col, field in enumerate(fields, 1):
        cell = ws.cell(row=1, column=col)
        cell.value = str(field.verbose_name) if field.verbose_name else field.name
        cell.font = Font(bold=True)
        cell.alignment = Alignment(horizontal='center')
    
    # 填充数据
    objects = model.objects.all()
    for row, obj in enumerate(objects, 2):
        for col, field in enumerate(fields, 1):
            # 处理外键字段
            if hasattr(field, 'related_model'):
                related_obj = getattr(obj, field.name)
                cell_value = str(related_obj) if related_obj else ""
            else:
                # 处理choice字段
                if hasattr(field, 'choices') and field.choices:
                    display_method = f"get_{field.name}_display"
                    if hasattr(obj, display_method):
                        cell_value = getattr(obj, display_method)()
                    else:
                        cell_value = str(getattr(obj, field.name)) if getattr(obj, field.name) is not None else ""
                else:
                    cell_value = str(getattr(obj, field.name)) if getattr(obj, field.name) is not None else ""
            
            ws.cell(row=row, column=col).value = cell_value
    
    # 调整列宽
    for col in range(1, len(fields) + 1):
        ws.column_dimensions[chr(64 + col)].width = 20
    
    # 保存文件
    wb.save(file_path)
    print(f"成功导出 {model.__name__} 数据到 {file_path}，共 {objects.count()} 条记录")


def main():
    # 检查参数
    if len(sys.argv) < 3:
        print("用法: python export_models.py <app_name.model_name> <export_file_path> [exclude_fields]")
        print("示例: python export_models.py house.MeterReading ./导出数据.xlsx id,created_at")
        sys.exit(1)
    
    # 解析参数
    model_path = sys.argv[1]
    export_path = sys.argv[2]
    exclude_fields = sys.argv[3].split(',') if len(sys.argv) > 3 else []
    
    # 分割应用名和模型名
    try:
        app_name, model_name = model_path.split('.')
    except ValueError:
        print("模型路径格式错误，请使用 <app_name.model_name> 格式")
        sys.exit(1)
    
    # 设置Django环境
    os.environ.setdefault("DJANGO_SETTINGS_MODULE", "house_crm.settings")
    django.setup()
    
    # 动态导入模型
    try:
        from django.apps import apps
        model = apps.get_model(app_name, model_name)
    except ImportError:
        print(f"无法导入应用 {app_name}")
        sys.exit(1)
    except LookupError:
        print(f"在应用 {app_name} 中找不到模型 {model_name}")
        sys.exit(1)
    
    # 导出数据
    export_model_to_excel(
        model=model,
        file_path=export_path,
        exclude_fields=exclude_fields
    )


if __name__ == "__main__":
    main()
